package gossip.v2.core;


import gossip.v2.GossipCore;
import gossip.v2.cluster.ExponentialBackoff;
import gossip.v2.message.*;
import gossip.v2.message.serializer.GossipMessageSerializer;
import gossip.v2.protocol.GossipHeartHandler;
import gossip.v2.protocol.GossipProtocol;
import gossip.v2.cluster.GossipNode;
import gossip.v2.protocol.GossipMessageProcessor;
import org.ricks.ioc.utils.Logger;
import org.ricks.net.transport.AioQuickClient;
import org.ricks.net.transport.AioSession;
import org.ricks.net.transport.WriteBuffer;
import java.io.IOException;
import java.net.ConnectException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

// 传输管理器
public class TransportManager {
    private final GossipCore core;
    private final Map<String, AioSession> sessions = new ConcurrentHashMap<>();
    private final GossipMessageSerializer serializer = new GossipMessageSerializer();
    private final Map<String, ReconnectContext> reconnectContexts = new ConcurrentHashMap<>();

    public TransportManager(GossipCore core) {
        this.core = core;
    }

    public synchronized void scheduleReconnect(GossipNode node) {
        if (core.isStopping()) return;
        if (node == null || node.equals(core.getLocalNode())) return;

        String nodeId = node.getNodeId();
        reconnectContexts.computeIfAbsent(nodeId, k -> new ReconnectContext(node, core)).scheduleNextAttempt();
    }

    // 优化方法：重连成功处理
    public void onReconnectSuccess(String nodeId, AioSession session) {
        ReconnectContext context = reconnectContexts.remove(nodeId);
        if (context != null) {
            context.shutdown();
        }

        // 注册会话并发送握手
        GossipNode node = core.getClusterManager().getNodeById(nodeId);
        if (node != null) {
            sessions.put(nodeId, session);
            core.getClusterManager().registerSession(node, session);
            context.sendHandshake(session); // 发送握手消息
        }

        Logger.info("{} - Successfully reconnected to: {}",core.getLocalNode().getNodeId(), nodeId);
    }

    // 新增方法：检查连接健康
    public boolean isConnectionHealthy(String nodeId) {
        AioSession session = sessions.get(nodeId);
        return session != null && !session.isInvalid();
    }

    public void onReconnectFailure(String nodeId) {
        // 失败时保持上下文继续重试
        Logger.debug("Reconnect attempt failed for: {}", nodeId);
    }


    // 新增方法：获取所有活动会话
    public Collection<AioSession> getAllSessions() {
        return new ArrayList<>(sessions.values());
    }

    public void onSessionClosed(AioSession session) {
        GossipNode node = core.getClusterManager().resolveNode(session);
        if (node != null) {
            sessions.remove(node.getNodeId());
            core.getClusterManager().clearNodeSession(session);
            core.getClusterManager().markAsSuspicious(node);

            // 触发重连逻辑
            if (core.getClusterManager().shouldReconnect(node)) {
                scheduleReconnect(node);
            }
        }
    }

    public void broadcast(GossipMessage message) {
        core.getClusterManager().getLiveNodes().values().stream()
                .filter(node -> !node.equals(core.getLocalNode()))
                .forEach(node -> send(node, message));
    }

    public void send(GossipNode node, GossipMessage message) {
        try {
            AioSession session = sessions.get(node.getNodeId());
            if (session == null || session.isInvalid()) {
                scheduleReconnect(node);
                return;
            }

            byte[] encoded = serializer.serialize(message);
            WriteBuffer writeBuffer = session.writeBuffer();
            writeBuffer.writeInt(encoded.length);
            writeBuffer.write(encoded);
            writeBuffer.flush();
        } catch (IOException e) {
            handleSendError(node, e);
        }
    }

    public void closeAllSessions() {
        sessions.values().forEach(session -> {
            try {
                if (!session.isInvalid()) {
                    GossipNode gossipNode = core.getClusterManager().resolveNode(session);
                    Logger.info("{} - 正在主动关闭节点 - {} close ", core.getLocalNode().getNodeId(), gossipNode.getNodeId());
                    session.close(true);
                }
            } catch (Exception e) {
                // Ignore
            }
        });
        sessions.clear();
    }

    public void shutdown() {
        closeAllSessions();
        reconnectContexts.values().forEach(ReconnectContext::shutdown);
        reconnectContexts.clear();
    }

    public synchronized void registerIncomingSession(GossipNode node, AioSession session) {
        sessions.put(node.getNodeId(), session);
        core.getClusterManager().registerSession(node, session);

//        // 立即发送握手响应
//        Map<String, GossipNode> knownNodes = core.getClusterManager().getAllKnownNodesWithAddress();
//        HandshakeMessage handshake = new HandshakeMessage(
//                core.getLocalNode(),
//                knownNodes,
//                HandshakeMessage.HandshakeType.RESPONSE
//        );
//        sendDirect(session, handshake);
    }

    public boolean hasActiveSession(GossipNode node) {
        return sessions.containsKey(node.getNodeId());
    }

    public synchronized void connectToNode(GossipNode node) {
        if (!shouldInitiateConnection(core.getLocalNode(), node)) {
            return;
        }
        if (node.equals(core.getLocalNode())) return;
        if (!core.getClusterManager().shouldReconnect(node)) return;

        String nodeId = node.getNodeId();
        if (sessions.containsKey(nodeId)) return;

        try {
            AioSession session = createSession(node);
            if (session != null) {
                sessions.put(nodeId, session);
                core.getClusterManager().registerSession(node, session);
                sendHandshake(session); // 发送握手请求
                Logger.info("{} - connect to {} success , {}  ", core.getLocalNode().getNodeId(), node.getNodeId(), session.hashCode());
            }
        } catch (IOException e) {
            handleConnectFailure(node, e);
        }
    }

    private void handleConnectFailure(GossipNode node, Exception e) {
        scheduleReconnect(node);
    }

    public boolean shouldInitiateConnection(GossipNode local, GossipNode remote) {
        // 确保节点间只有一条连接
        return local.getNodeId().hashCode() < remote.getNodeId().hashCode();
    }

    public AioSession createSession(GossipNode node) throws IOException {
        try {
            GossipMessageProcessor processor = new GossipMessageProcessor(core);
            AioQuickClient client = new AioQuickClient(
                    node.getIp(),
                    node.getPort(),
                    new GossipProtocol(core),
                    processor
            );
            // 添加心跳插件
            processor.addPlugin(new GossipHeartHandler(2000, 5000, TimeUnit.MILLISECONDS, core));
            client.connectTimeout(3000);
            return client.start();
        } catch (ConnectException e) {
            return null;
        }
    }

    public void sendHeartbeat(AioSession session, HeartbeatMessage.HeartbeatType heartbeatType) {
        sendDirect(session, new HeartbeatMessage(core.getLocalNode(), System.currentTimeMillis(),heartbeatType));
//        Logger.info("{} - {} 正在发送心跳 ",core.getLocalNode().getNodeId(), session.hashCode());
    }

    private void sendHandshake(AioSession session) {
        // 立即发送握手响应
        Map<String, GossipNode> knownNodes = core.getClusterManager().getAllKnownNodesWithAddress();
        HandshakeMessage handshake = new HandshakeMessage(
                core.getLocalNode(),
                knownNodes,
                HandshakeMessage.HandshakeType.REQUEST
        );
        sendDirect(session, handshake);
    }

    public void sendDirect(AioSession session, GossipMessage message) {
        try {
            if(session.isInvalid()) {
                Logger.error("session:{} - {} 失效,无法写出", null,session.hashCode(), message.getSenderId());
            }
            byte[] data = serializer.serialize(message);
            WriteBuffer writeBuffer = session.writeBuffer();
            writeBuffer.writeInt(data.length);
            writeBuffer.write(data);
            writeBuffer.flush();
        } catch (IOException e) {
            handleSendError(session, e);
        }
    }

    private void handleSendError(GossipNode node, Exception e) {
        Logger.info(" 触发了发送失败主动关闭");
        sessions.remove(node.getNodeId());
        if (core.getClusterManager().isSeedNode(node.getNodeId())) {
            scheduleReconnect(node);
        } else {
            core.getClusterManager().markAsSuspicious(node);
        }
    }

    private void handleSendError(AioSession session, Exception e) {
        Logger.info(" 触发了发送失败主动关闭");
        GossipNode node = core.getClusterManager().resolveNode(session);
        if (node != null) {
            sessions.remove(node.getNodeId());
            core.getClusterManager().markAsSuspicious(node);
        }
        try {
            session.close(true);
            Logger.info("{} - 触发了发送失败{} - {} 主动关闭",core.getLocalNode().getNodeId(),node.getNodeId(),session );
        } catch (Exception ex) {
            // Ignore
        }
    }
}
